package cn.edu.hitsz.compiler.asm;

import cn.edu.hitsz.compiler.NotImplementedException;
import cn.edu.hitsz.compiler.ir.*;
import cn.edu.hitsz.compiler.utils.FileUtils;
import org.jetbrains.annotations.NotNull;

import java.util.*;


/**
 * TODO: 实验四: 实现汇编生成
 * <br>
 * 在编译器的整体框架中, 代码生成可以称作后端, 而前面的所有工作都可称为前端.
 * <br>
 * 在前端完成的所有工作中, 都是与目标平台无关的, 而后端的工作为将前端生成的目标平台无关信息
 * 根据目标平台生成汇编代码. 前后端的分离有利于实现编译器面向不同平台生成汇编代码. 由于前后
 * 端分离的原因, 有可能前端生成的中间代码并不符合目标平台的汇编代码特点. 具体到本项目你可以
 * 尝试加入一个方法将中间代码调整为更接近 risc-v 汇编的形式, 这样会有利于汇编代码的生成.
 * <br>
 * 为保证实现上的自由, 框架中并未对后端提供基建, 在具体实现时可自行设计相关数据结构.
 *
 * @see AssemblyGenerator#run() 代码生成与寄存器分配
 */
public class AssemblyGenerator {

    List<Instruction> instructions = new ArrayList<>();
    //内存分配
    Map<IRVariable, Integer> memMap = IRVariable.getMemMap();
    //寄存器分配
    Map<IRVariable, Integer> regMap = new HashMap<>();
    Set<DUchain> spillSet = new HashSet<>();

    List<String> asmCodes = new LinkedList<>();

    /**
     * 加载前端提供的中间代码
     * <br>
     * 视具体实现而定, 在加载中或加载后会生成一些在代码生成中会用到的信息. 如变量的引用
     * 信息. 这些信息可以通过简单的映射维护, 或者自行增加记录信息的数据结构.
     *
     * @param originInstructions 前端提供的中间代码
     */
    public void loadIR(List<Instruction> originInstructions) {
        instructions = originInstructions;
    }

    class DUchain
    {
        int def;
        int use;
        IRVariable variable;
        int regid;
        int usecnt;
        boolean spilled;
        DUchain(IRVariable variable, int def)
        {
            this.def = def;
            this.use = def;
            this.variable = variable;
            usecnt = 1;
            spilled = false;
            regid = -1;
        }

        int weight()
        {
            return (use - def) / usecnt;
        }
    }

    /**
     * 生成def-use区间
     * @return def-use区间列表
     */
    private List<DUchain> generateDuchain()
    {
        List<DUchain> duchainList = new ArrayList<>();
        Map<IRVariable, DUchain> duchainMap = new HashMap<>();

        //生成活跃区间
        for(int i=0;i<instructions.size();i++)
        {
            Instruction instruction = instructions.get(i);
            switch (instruction.getKind())
            {
                case MOV ->
                {
                    IRVariable result = instruction.getResult();
                    IRValue from = instruction.getFrom();
                    if (from.isIRVariable())
                    {
                        DUchain temp = duchainMap.get((IRVariable) from);
                        temp.use = i;
                        temp.usecnt++;
                    }
                    DUchain du = new DUchain(result, i);
                    if (i < instructions.size()-1 && instructions.get(i+1).getKind() == InstructionKind.STORE)
                        du.spilled = true;
                    duchainList.add(du);
                    if (duchainMap.containsKey(result))
                        duchainMap.remove(result);
                    duchainMap.put(result, du);
                }
                case ADD, SUB, MUL ->
                {
                    IRValue lhs = instruction.getLHS();
                    IRValue rhs = instruction.getRHS();
                    IRVariable result = instruction.getResult();
                    if (lhs.isIRVariable())
                    {
                        DUchain temp = duchainMap.get((IRVariable) lhs);
                        temp.use = i;
                        temp.usecnt++;
                    }
                    if (rhs.isIRVariable())
                    {
                        DUchain temp = duchainMap.get((IRVariable) rhs);
                        temp.use = i;
                        temp.usecnt++;
                    }
                    DUchain du = new DUchain(result, i);
                    if (i < instructions.size()-1 && instructions.get(i+1).getKind() == InstructionKind.STORE)
                        du.spilled = true;
                    duchainList.add(du);
                    if (duchainMap.containsKey(result))
                        duchainMap.remove(result);
                    duchainMap.put(result, du);
                }
                case LOAD ->
                {
                    IRValue lhs = instruction.getLHS();
                    DUchain du = new DUchain((IRVariable) lhs, i);
                    du.spilled = true;
                    duchainList.add(du);
                    if (duchainMap.containsKey((IRVariable) lhs))
                        duchainMap.remove((IRVariable) lhs);
                    duchainMap.put((IRVariable) lhs, du);
                }
                case STORE ->
                {
                    IRValue lhs = instruction.getLHS();
                    DUchain temp = duchainMap.get((IRVariable) lhs);
                    temp.use = i;
                    temp.usecnt++;
                }
                case RET ->
                {
                    IRValue ret = instruction.getReturnValue();
                    if (ret.isIRVariable())
                    {
                        DUchain temp = duchainMap.get((IRVariable) ret);
                        temp.use = i;
                        temp.usecnt++;
                    }
                }
            }
        }
        return duchainList;
    }

    /**
     * 线性扫描法进行寄存器分配
     * @return 新的spill集合
     */
    private @NotNull Set<DUchain> Reg_allocated(List<DUchain> duchainList)
    {
        int available = 0x3f;//t6备用
        Set<DUchain> spilled = new HashSet<>();
        Queue<DUchain> heap = new PriorityQueue<>(duchainList.size() , new Comparator<DUchain>() {
            @Override
            public int compare(DUchain o1, DUchain o2) {
                return o2.weight() - o1.weight();
            }
        });
        Queue<DUchain> end = new PriorityQueue<>(duchainList.size() , new Comparator<DUchain>() {
            @Override
            public int compare(DUchain o1, DUchain o2) {
                return o1.use - o2.use;
            }
        });

        int j = 0;
        Map<IRVariable, DUchain> duchainMap = new HashMap<>();

        for(int i=0;i<instructions.size();i++)
        {
            Instruction instruction = instructions.get(i);

            //回收不再使用的变量的寄存器
            while (!end.isEmpty())
            {
                DUchain temp = end.poll();
                if(temp.regid > -1 && temp.use <= i)
                    available |= 1 << temp.regid;
                else if (temp.use > i)
                {
                    end.add(temp);
                    break;
                }
            }

            DUchain now;
            if (j < duchainList.size())
                now = duchainList.get(j);
            else
                now = duchainList.get(duchainList.size()-1);

            switch (instruction.getKind())
            {
                case MOV, ADD, SUB, MUL, LOAD ->
                {
                    j++;
                    if (!now.spilled)
                        heap.add(now);
                    else
                        spilled.add(now);
                    end.add(now);
                    if (duchainMap.containsKey(now.variable))
                        duchainMap.remove(now.variable);
                    duchainMap.put(now.variable, now);

                    //寄存器均被分配
                    if (available == 0)
                    {
                        //spill use频率最小即weight最大的区间
                        while (!heap.isEmpty())
                        {
                            DUchain temp = heap.poll();
                            if (temp.use > i && temp.regid > -1)
                            {
                                spilled.add(temp);
                                //回收寄存器
                                available |= 1 << temp.regid;
                                temp.regid = -1;
                                break;
                            }
                        }
                    }
                    //分配寄存器
                    if (available == 0)
                    {
                        spilled.add(now);
                        continue;
                    }

                    for(int k=0;k<6;k++)
                    {
                        if (((available>>k)&1) == 1)
                        {
                            available = available ^ (1<<k);
                            now.regid = k;
                            break;
                        }
                    }

                }
                case STORE, RET -> {}
            }

        }
        return spilled;
    }

    /**
     * 生成spill后的中间代码
     * @return 调整后的中间代码
     */
    private List<Instruction> spill(List<DUchain> dUchainList)
    {
        List<Instruction> temp = new ArrayList<>();
        int i = 0;
        Map<IRVariable, DUchain> duchainMap = new HashMap<>();

        for(Instruction instruction : instructions)
        {
            DUchain now;
            if (i < dUchainList.size())
                now = dUchainList.get(i);
            else
                now = dUchainList.get(dUchainList.size()-1);
            switch (instruction.getKind())
            {
                case MOV, ADD, SUB, MUL, LOAD ->
                {
                    i++;
                    if (duchainMap.containsKey(now.variable))
                        duchainMap.remove(now.variable);
                    duchainMap.put(now.variable, now);
                }
                case STORE, RET -> {}
            }
            //spill
            switch (instruction.getKind())
            {
                case MOV ->
                {
                    IRVariable result = instruction.getResult();
                    IRValue from = instruction.getFrom();
                    if (from.isIRVariable() && spillSet.contains(duchainMap.get((IRVariable) from))) //在 use 前加 LOAD
                        temp.add(Instruction.createLoad((IRVariable) from, IRImmediate.of(memMap.get((IRVariable) from))));
                    temp.add(instruction);
                    if (spillSet.contains(duchainMap.get(result))) //在 def 后加 STORE
                        temp.add(Instruction.createStore(result, IRImmediate.of(memMap.get(result))));
                }
                case ADD, SUB, MUL ->
                {
                    IRValue lhs = instruction.getLHS();
                    IRValue rhs = instruction.getRHS();
                    IRVariable result = instruction.getResult();
                    if (lhs.isIRVariable() && spillSet.contains(duchainMap.get((IRVariable) lhs))) //在 use 前加 LOAD
                        temp.add(Instruction.createLoad((IRVariable) lhs, IRImmediate.of(memMap.get((IRVariable) lhs))));
                    if (rhs.isIRVariable() && spillSet.contains(duchainMap.get((IRVariable) rhs))) //在 use 前加 LOAD
                        temp.add(Instruction.createLoad((IRVariable) rhs, IRImmediate.of(memMap.get((IRVariable) rhs))));
                    temp.add(instruction);
                    if (spillSet.contains(duchainMap.get(result))) //在 def 后加 STORE
                        temp.add(Instruction.createStore(result, IRImmediate.of(memMap.get(result))));
                }
                case LOAD, STORE -> {}
                case RET->
                {
                    IRValue ret = instruction.getReturnValue();
                    if (ret.isIRVariable() && spillSet.contains(duchainMap.get((IRVariable) ret))) //在 use 前加 LOAD
                    {
                        temp.add(Instruction.createLoad((IRVariable) ret, IRImmediate.of(memMap.get((IRVariable) ret))));
                    }
                    temp.add(instruction);
                }
            }
        }
        return temp;
    }

    /**
     * 生成汇编代码
     */
    private void generateAsmCodes(List<DUchain> duchainList)
    {
        asmCodes.add(".text");
        Map<IRVariable, DUchain> duchainMap = new HashMap<>();
        int i = 0;
        for (Instruction instruction : instructions)
        {
            DUchain now;
            if (i < duchainList.size())
                now = duchainList.get(i);
            else
                now = duchainList.get(duchainList.size()-1);

            switch (instruction.getKind())
            {
                case MOV, ADD, SUB, MUL, LOAD ->
                {
                    i++;
                    if (duchainMap.containsKey(now.variable))
                        duchainMap.remove(now.variable);
                    duchainMap.put(now.variable, now);
                    //记录分配寄存器情况
                    if (regMap.containsKey(now.variable))
                        regMap.remove(now.variable);
                    regMap.put(now.variable, now.regid);
                }
                case STORE, RET -> {}
            }


            switch (instruction.getKind())
            {
                case ADD ->
                {
                    IRValue lhs = instruction.getLHS();
                    IRValue rhs = instruction.getRHS();
                    IRVariable result = instruction.getResult();
                    int rd = regMap.get(result);
                    if (lhs.isImmediate() && rhs.isImmediate())
                    {
                        IRImmediate temp1 = (IRImmediate) lhs;
                        IRImmediate temp2 = (IRImmediate) rhs;
                        int res = temp1.getValue() + temp2.getValue();
                        asmCodes.add(AsmCode.li(rd, res));
                    }
                    else if (lhs.isImmediate())
                    {
                        IRImmediate temp1 = (IRImmediate) lhs;
                        asmCodes.addAll(AsmCode.addi(rd, regMap.get((IRVariable) rhs), temp1.getValue()));
                    }
                    else if (rhs.isImmediate())
                    {
                        IRImmediate temp2 = (IRImmediate) rhs;
                        asmCodes.addAll(AsmCode.addi(rd, regMap.get((IRVariable) lhs), temp2.getValue()));
                    }
                    else
                        asmCodes.add(AsmCode.add(rd, regMap.get((IRVariable) lhs), regMap.get((IRVariable) rhs)));
                }
                case SUB ->
                {
                    IRValue lhs = instruction.getLHS();
                    IRValue rhs = instruction.getRHS();
                    IRVariable result = instruction.getResult();
                    int rd = regMap.get(result);
                    if (lhs.isImmediate() && rhs.isImmediate())
                    {
                        IRImmediate temp1 = (IRImmediate) lhs;
                        IRImmediate temp2 = (IRImmediate) rhs;
                        int res = temp1.getValue() - temp2.getValue();
                        asmCodes.add(AsmCode.li(rd, res));
                    }
                    else if (lhs.isImmediate())
                    {
                        IRImmediate temp1 = (IRImmediate) lhs;
                        asmCodes.add(AsmCode.li(rd, temp1.getValue()));//addi x0,temp1.getValue
                        asmCodes.add(AsmCode.sub(rd, rd, regMap.get((IRVariable) rhs)));
                    }
                    else if (rhs.isImmediate())
                    {
                        IRImmediate temp2 = (IRImmediate) rhs;
                        asmCodes.addAll(AsmCode.addi(rd, regMap.get((IRVariable) lhs), -temp2.getValue()));//addi rd,-temp2.getValue 即subi
                    }
                    else
                        asmCodes.add(AsmCode.sub(rd, regMap.get((IRVariable) lhs), regMap.get((IRVariable) rhs)));
                }
                case MUL ->
                {
                    IRValue lhs = instruction.getLHS();
                    IRValue rhs = instruction.getRHS();
                    IRVariable result = instruction.getResult();
                    int rd = regMap.get(result);
                    if (lhs.isImmediate() && rhs.isImmediate())
                    {
                        IRImmediate temp1 = (IRImmediate) lhs;
                        IRImmediate temp2 = (IRImmediate) rhs;
                        int res = temp1.getValue() * temp2.getValue();
                        asmCodes.add(AsmCode.li(rd, res));
                    }
                    else if (lhs.isImmediate())
                    {
                        IRImmediate temp1 = (IRImmediate) lhs;
                        //没有muli
                        asmCodes.add(AsmCode.li(rd, temp1.getValue()));
                        asmCodes.add(AsmCode.mul(rd, rd, regMap.get((IRVariable) rhs)));
                    }
                    else if (rhs.isImmediate())
                    {
                        IRImmediate temp2 = (IRImmediate) rhs;
                        //没有muli
                        asmCodes.add(AsmCode.li(rd, temp2.getValue()));
                        asmCodes.add(AsmCode.mul(rd, rd, regMap.get((IRVariable) lhs)));
                    }
                    else
                        asmCodes.add(AsmCode.mul(rd, regMap.get((IRVariable) lhs), regMap.get((IRVariable) rhs)));
                }
                case MOV ->
                {
                    IRVariable result = instruction.getResult();
                    IRValue from = instruction.getFrom();
                    int rd = regMap.get(result);
                    int rs;
                    if (from.isImmediate())
                    {
                        IRImmediate temp = (IRImmediate) from;
                        asmCodes.add(AsmCode.li(rd, temp.getValue()));
                    }
                    else if (rd != (rs=regMap.get((IRVariable) from)))
                        asmCodes.add(AsmCode.mv(rd, rs));
                }
                case LOAD ->
                {
                    IRVariable lhs = (IRVariable) instruction.getLHS();
                    IRImmediate rhs = (IRImmediate )instruction.getRHS();
                    asmCodes.addAll(AsmCode.lw(regMap.get(lhs), rhs.getValue()));
                }
                case STORE ->
                {
                    IRVariable lhs = (IRVariable) instruction.getLHS();
                    IRImmediate rhs = (IRImmediate )instruction.getRHS();
                    asmCodes.addAll(AsmCode.sw(regMap.get(lhs), rhs.getValue()));
                }
                case RET ->
                {
                    IRValue ret = instruction.getReturnValue();
                    if (ret.isImmediate())
                    {
                        IRImmediate res = (IRImmediate) ret;
                        asmCodes.add(AsmCode.li(7, res.getValue()));// li a0, res.getValue()
                    }
                    else
                        asmCodes.add(AsmCode.mv(7, regMap.get((IRVariable) ret)));// mv a0, res.getValue()
                }
            }
            asmCodes.add("#" + instruction.toString());
        }
    }

    /**
     * 执行代码生成.
     * <br>
     * 根据理论课的做法, 在代码生成时同时完成寄存器分配的工作. 若你觉得这样的做法不好,
     * 也可以将寄存器分配和代码生成分开进行.
     * <br>
     * 提示: 寄存器分配中需要的信息较多, 关于全局的与代码生成过程无关的信息建议在代码生
     * 成前完成建立, 与代码生成的过程相关的信息可自行设计数据结构进行记录并动态维护.
     */
    public void run() {

        List<DUchain> duchainList;
        while(true)
        {
            duchainList = generateDuchain();
            spillSet = Reg_allocated(duchainList);
            int cnt = instructions.size();
            instructions = spill(duchainList);
            if (cnt == instructions.size()) //说明没有新的spill操作
                break;
        }
        FileUtils.writeLines("data/out/new_intermediate_code.txt", instructions.stream().map(Instruction::toString).toList());
        generateAsmCodes(duchainList);

    }


    /**
     * 输出汇编代码到文件
     *
     * @param path 输出文件路径
     */
    public void dump(String path) {
        FileUtils.writeLines(path, asmCodes);
    }
}

